##############################################################################################
#
#
#  A code for the recursions of Shasha-Zhang (1989), Klein (1998), and
#  Demain et al. (2007).
#  These algorithms are implemented with space complexity = time complexity.
#  It is possible to reduce the space complexity to O(mn) by impementing acording to section 5 of
#  http://www.mit.edu/~oweimann/Publications/TEDinTALG.pdf.
#
#  To run it, first create two Forest instances:
#  F = Forest(sigma="your alphabet")   #(default is "01")
#  To create a random forest use F.create_random_tree(size), otherwise create your own forests.
#  And same for G.
#
#  Then create a Shasha/Klein/Demain instance by say:
#
#  s = Shasha()
#  k =Klein2()
#  d = DMRW()
#
#  And calc the TED (tree edit distance) by:
#  s.ted(F,G) or k.ted(F,G) or d.ted(F,G)
#
##############################################################################################

from random import choice

class Node(object):
    "node in an ordered tree"

    def __init__(self,label,children=None,idn=None,parent=None):
        self.label = label
        self.id = idn
        self.parent = parent
        self.nbrs=[]
        self.n=1
        if children:
            self.nbrs.extend(children)
            for c in children:
                self.n = self.n + c.n

    def __str__(self):
        if len(self.nbrs)>0:
            return str(self.id)+":"+str(self.nbrs[0].id)+"."+str(self.nbrs[-1].id)
        else:
            return str(self.id)+":"+"None"
        
    def depth_first(self):
        for n in self.nbrs:
            for node in n.depth_first():
                yield node
        yield self


class Forest(object):
    "class for labeled rooted ordered forest"

    def __init__(self,roots=[],n=0,sigma="01"):
        self.roots = roots
        self.sigma = sigma
        self.n=n
        if not self.empty():
            tmp = [x.n for x in self.roots]
            self.heavy = tmp.index(max(tmp))

    def free(self):
        for x in self.depth_first():
                del x.parent
                del x.nbrs
        del self.roots
        
    def empty(self):
        return 0==len(self.roots)

    def balanced(self,n):
        nodes = map(lambda x: Node(choice(self.sigma),idn=x), range(pow(2,n)-1))
        for i in range(pow(2,n-1)-1):
            nodes[i].nbrs = [nodes[2*i+1],nodes[2*i+2]]
        for i in range(1,pow(2,n)-1):
            nodes[i].parent = nodes[(i-1)/2]
        self.roots = [nodes[0]]
        self.recalc_n()

                                   
    def zigzag(self,n):
        zig = Node(choice(self.sigma),idn=0)
        for i in range (1,n,4):
            tmp = Node(choice(self.sigma),idn=i)
            zag = Node(choice(self.sigma),children=[tmp,zig],idn=i+1)
            tmp.parent=zag
            zig.parent=zag
            tmp = Node(choice(self.sigma),idn=i+2)
            zig = Node(choice(self.sigma),children=[zag,tmp],idn=i+3)
            tmp.parent=zig
            zag.parent=zig
        self.roots = [zig]
        self.recalc_n()
            
        
    def create_random_tree(self,n):
        self.roots = [Node(choice(self.sigma),idn=0)]
        for i in range(1,n):
            node = self.roots[0]
            while True:
                d = choice([0,1])
                try:
                    node.n = node.n+1
                    node = node.nbrs[d]
                except IndexError:
                        node.nbrs.append(Node(choice(self.sigma),idn=i,parent=node))
                        break
            self.n=n
    
    def del_left(self):
        if self.empty():
            raise IndexError
        r = self.roots[0].nbrs[:]
        r.extend(self.roots[1:])
        return Forest(r,self.n-1,self.sigma)
        
    def del_right(self):
        if self.empty():
            raise IndexError
        r = self.roots[:-1]
        r.extend(self.roots[-1].nbrs[:])
        return Forest(r,self.n-1,self.sigma)

    def match_left(self):
        if self.empty():
            raise IndexError
        r = self.roots[0].nbrs[:]
        return (Forest(self.roots[1:],self.n - self.roots[0].n,self.sigma) , Forest(r,self.roots[0].n -1,self.sigma))

    def match_right(self):
        if self.empty():
            raise IndexError
        r = self.roots[-1].nbrs[:]
        return (Forest(self.roots[:-1],self.n - self.roots[-1].n,self.sigma) , Forest(r,self.roots[-1].n -1,self.sigma))

    def depth_first(self):
        if not self.empty():
            for r in self.roots:
                for node in r.depth_first():
                    yield node

    def __str__(self):
       s = "("
       for r in self.roots:
           s = s + "," + str(r)
       return s+")"
   
    def calc_heavy(self):
        if self.empty():
            return

        tmp = [x.n for x in self.roots]
        self.heavy = tmp.index(max(tmp))
        
        for i in self.depth_first():
            if len(i.nbrs)>0:
                tmp = [x.n for x in i.nbrs]
                i.heavy = tmp.index(max(tmp))

    def toplight(self):
        if self.empty() or self.n==1:
            return []
        res = self.roots[0].nbrs[:]
        h = self.roots[0].heavy
        res.pop(h)
        res = [Forest([x],x.n,self.sigma) for x in res]
        f = Forest([self.roots[0].nbrs[h]],self.roots[0].nbrs[h].n,self.sigma)
        res.extend(f.toplight())
        return res

    def change_root(self,root, left):

        #special handling when root is not changed
        if root == self.roots[0]:
            i=root.nbrs.index(left)
            nbrs = root.nbrs[i:] + root.nbrs[:i]
            new_root = Node(root.label,nbrs,root.id,None)
#            print "created new root:",new_root
            F = Forest([new_root],self.n,self.sigma)
            F.recalc_n()
            return F            

        nbrs = []
        p = root.parent
        if p==left:
            nbrs = [left] + root.nbrs
        else:
            i=root.nbrs.index(left)
            nbrs = root.nbrs[i:] + [p] + root.nbrs[:i]

        new_root = Node(root.label,nbrs,root.id,None)
#        print "created new root:",new_root
        
        prev = root
        new_prev = new_root
        crnt = p

        while True:
            p = crnt.parent
            i = crnt.nbrs.index(prev)
            if p is None:
                nbrs = crnt.nbrs[i+1:] + crnt.nbrs[:i]
            else:
                nbrs = crnt.nbrs[i+1:] + [p] + crnt.nbrs[:i]
            new = Node(crnt.label,nbrs,crnt.id,new_prev)
#            print "created new node:",new

            #update new node in parent
            i = new_prev.nbrs.index(crnt)
            new_prev.nbrs[i] = new
            
            if crnt==self.roots[0]:
                break
            else:
                prev = crnt
                new_prev = new
                crnt = p

        F = Forest([new_root],self.n,self.sigma)
        F.recalc_n()
        return F

    def recalc_n(self):
        for x in self.depth_first():
            if len(x.nbrs)==0:
                x.n = 1
            else:
                x.n = 1 + sum(y.n for y in x.nbrs)
        self.n = self.roots[0].n
        return self.n
            

class Shasha:

    def __init__(self):
        self.t={}

    def ted(self,F,G):
        #print str(F)
        #print str(G)
        #print "*************"
        if F.empty() and G.empty():
            self.t[(str(F),str(G))] = 0

        elif F.empty():
            key = (str(F),str(G.del_right()))
            if not self.t.has_key(key):
                self.ted(F,G.del_right())
            self.t[(str(F),str(G))] = self.t[key] + 1

        elif G.empty():
            key = (str(F.del_right()),str(G))
            if not self.t.has_key(key):
                self.ted(F.del_right(),G)
            self.t[(str(F),str(G))] = self.t[key] + 1

        else:
            key1 = (str(F),str(G.del_right()))
            if not self.t.has_key(key1):
                self.ted(F,G.del_right())
            key2 = (str(F.del_right()),str(G))
            if not self.t.has_key(key2):
                self.ted(F.del_right(),G)
            (F3,F4) = F.match_right()
            (G3,G4) = G.match_right()
            print str(F)
            print str(F3),str(G3)
            print str(F4),str(G4)
            print ""
            key3 = (str(F3),str(G3))
            key4 = (str(F4),str(G4))
            if not self.t.has_key(key3):
                self.ted(F3,G3)
            if not self.t.has_key(key4):
                self.ted(F4,G4)

            self.t[(str(F),str(G))] = min(self.t[key1] + 1 , self.t[key2] + 1 , self.t[key3] + self.t[key4] + (1 - (F.roots[-1].label==G.roots[-1].label)))

            return self.t[str(F),str(G)]

class Klein:

    def __init__(self):
        self.t={}

    def ted(self,F,G):
        if F.n >= G.n:
            self.master=0
            F.calc_heavy()
        else:
            self.master=1
            G.calc_heavy()
        return self.ted1(F,G)
        
    def ted1(self,F,G):

        if F.empty() and G.empty():
            self.t[(str(F),str(G))] = 0

        elif F.empty():
            if self.master==0 or (self.master==1 and G.heavy==0 and len(G.roots)>1):
                key = (str(F),str(G.del_right()))
                if not self.t.has_key(key):
                    self.ted1(F,G.del_right())
            else:
                key = (str(F),str(G.del_left()))
                if not self.t.has_key(key):
                    self.ted1(F,G.del_left())
            self.t[(str(F),str(G))] = self.t[key] + 1

        elif G.empty():
            if self.master==1 or (self.master==0 and F.heavy==0 and len(F.roots)>1):
                key = (str(F.del_right()),str(G))
                if not self.t.has_key(key):
                    self.ted1(F.del_right(),G)
            else:
                key = (str(F.del_left()),str(G))
                if not self.t.has_key(key):
                    self.ted1(F.del_left(),G)                        
            self.t[(str(F),str(G))] = self.t[key] + 1

        else:
            if (self.master==0 and F.heavy==0 and len(F.roots)>1) or (self.master==1 and G.heavy==0 and len(G.roots)>1):
                key1 = (str(F),str(G.del_right()))
                if not self.t.has_key(key1):
                    self.ted1(F,G.del_right())
                key2 = (str(F.del_right()),str(G))
                if not self.t.has_key(key2):
                    self.ted1(F.del_right(),G)
                (F3,F4) = F.match_right()
                (G3,G4) = G.match_right()
                key3 = (str(F3),str(G3))
                key4 = (str(F4),str(G4))
                if not self.t.has_key(key3):
                    self.ted1(F3,G3)
                if not self.t.has_key(key4):
                    self.ted1(F4,G4)

                self.t[(str(F),str(G))] = min(self.t[key1] + 1 , self.t[key2] + 1 , self.t[key3] + self.t[key4] + (1-(F.roots[-1].label==G.roots[-1].label)))

            else:

                key1 = (str(F),str(G.del_left()))
                if not self.t.has_key(key1):
                    self.ted1(F,G.del_left())
                key2 = (str(F.del_left()),str(G))
                if not self.t.has_key(key2):
                    self.ted1(F.del_left(),G)
                (F3,F4) = F.match_left()
                (G3,G4) = G.match_left()
                key3 = (str(F3),str(G3))
                key4 = (str(F4),str(G4))
                if not self.t.has_key(key3):
                    self.ted1(F3,G3)
                if not self.t.has_key(key4):
                    self.ted1(F4,G4)

                self.t[(str(F),str(G))] = min(self.t[key1] + 1 , self.t[key2] + 1 , self.t[key3] + self.t[key4] + (1-F.roots[0].label==G.roots[0].label))

            return self.t[str(F),str(G)]


class Klein2:

    def __init__(self):
        self.t={}

    def ted(self,F,G):
        if F.n >= G.n:
            self.master=0
            F.calc_heavy()
        else:
            self.master=1
            G.calc_heavy()
        return self.ted2(F,G)

    def ted2(self,F,G):
        if (self.master==0):
            tl = F.toplight()
            for t in tl:
                self.ted2(t,G)
        else:
            tl = G.toplight()
            for t in tl:
                self.ted2(F,t)
        return self.ted3(F,G)
        
    def ted3(self,F,G):

        if F.empty() and G.empty():
            self.t[(str(F),str(G))] = 0

        elif F.empty():
            if self.master==0:
                keyr = (str(F),str(G.del_right()))
                if not self.t.has_key(keyr):
                    self.ted3(F,G.del_right())
                keyl = (str(F),str(G.del_left()))
                if not self.t.has_key(keyl):
                    self.ted3(F,G.del_left())
                self.t[(str(F),str(G))] = min(self.t[keyr], self.t[keyl]) + 1
                
            elif self.master==1 and G.heavy==0 and len(G.roots)>1:
                key = (str(F),str(G.del_right()))
                if not self.t.has_key(key):
                    self.ted3(F,G.del_right())
                self.t[(str(F),str(G))] = self.t[key] + 1
            else:
                key = (str(F),str(G.del_left()))
                if not self.t.has_key(key):
                    self.ted3(F,G.del_left())
                self.t[(str(F),str(G))] = self.t[key] + 1

        elif G.empty():
            if self.master==1:
                keyr = (str(F.del_right()),str(G))
                if not self.t.has_key(keyr):
                    self.ted3(F.del_right(),G)
                keyl = (str(F.del_left()),str(G))
                if not self.t.has_key(keyl):
                    self.ted3(F.del_left(),G)
                self.t[(str(F),str(G))] = min(self.t[keyr], self.t[keyl]) + 1

            elif self.master==0 and F.heavy==0 and len(F.roots)>1:
                key = (str(F.del_right()),str(G))
                if not self.t.has_key(key):
                    self.ted3(F.del_right(),G)
                self.t[(str(F),str(G))] = self.t[key] + 1
            else:
                key = (str(F.del_left()),str(G))
                if not self.t.has_key(key):
                    self.ted3(F.del_left(),G)                        
                self.t[(str(F),str(G))] = self.t[key] + 1

        else:

            if (self.master==0 and F.heavy==0 and len(F.roots)>1) or (self.master==1 and G.heavy==0 and len(G.roots)>1):
                if self.master==0:                    
                    key0 = (str(F),str(G.del_left()))
                    if not self.t.has_key(key0):
                        self.ted3(F,G.del_left())
                else:
                    key0 = (str(F.del_left()),str(G))
                    if not self.t.has_key(key0):
                        self.ted3(F.del_left(),G)

                key1 = (str(F),str(G.del_right()))
                if not self.t.has_key(key1):
                    self.ted3(F,G.del_right())
                key2 = (str(F.del_right()),str(G))
                if not self.t.has_key(key2):
                    self.ted3(F.del_right(),G)
                (F3,F4) = F.match_right()
                (G3,G4) = G.match_right()
                key3 = (str(F3),str(G3))
                key4 = (str(F4),str(G4))
                if not self.t.has_key(key3):
                    self.ted3(F3,G3)
                if not self.t.has_key(key4):
                    print "problems!!!"
                    self.ted3(F4,G4)

                self.t[(str(F),str(G))] = min(self.t[key0] + 1 ,
                                              self.t[key1] + 1 ,
                                              self.t[key2] + 1 ,
                                              self.t[key3] + self.t[key4] + (1-(F.roots[-1].label==G.roots[-1].label)))

            else:
                
                if self.master==0:                    
                    key0 = (str(F),str(G.del_right()))
                    if not self.t.has_key(key0):
                        self.ted3(F,G.del_right())
                else:
                    key0 = (str(F.del_right()),str(G))
                    if not self.t.has_key(key0):
                        self.ted3(F.del_right(),G)

                key1 = (str(F),str(G.del_left()))
                if not self.t.has_key(key1):
                    self.ted3(F,G.del_left())
                key2 = (str(F.del_left()),str(G))
                if not self.t.has_key(key2):
                    self.ted3(F.del_left(),G)
                (F3,F4) = F.match_left()
                (G3,G4) = G.match_left()
                key3 = (str(F3),str(G3))
                key4 = (str(F4),str(G4))
                if not self.t.has_key(key3):
                    self.ted3(F3,G3)
                if not self.t.has_key(key4):
                    print "problems!!!!"
                    self.ted3(F4,G4)

                self.t[(str(F),str(G))] = min(self.t[key0] + 1 ,
                                              self.t[key1] + 1 ,
                                              self.t[key2] + 1 ,
                                              self.t[key3] + self.t[key4] + (1-(F.roots[0].label==G.roots[0].label)))

            return self.t[str(F),str(G)]

class DMRW:

    def __init__(self):
        self.t={}

    def ted(self,F,G):
        F.calc_heavy()
        G.calc_heavy()
        return self.ted2(F,G)
    
    def ted2(self,F,G):
#        print "now calculating",(str(F),str(G)),F.n,G.n
        if F.n >= G.n:
            master=0
        else:
            master=1

        if (master==0):
            tl = F.toplight()
            for t in tl:
                self.ted2(t,G)
        else:
            tl = G.toplight()
            for t in tl:
                self.ted2(F,t)
        return self.ted3(F,G,master)
        
    def ted3(self,F,G,master):
#        print "in step3",(str(F),str(G)), "master:",master      

        if F.empty() and G.empty():
            self.t[(str(F),str(G))] = 0

        elif F.empty():
            if master==0:
                keyr = (str(F),str(G.del_right()))
                if not self.t.has_key(keyr):
                    self.ted3(F,G.del_right(),master)
                keyl = (str(F),str(G.del_left()))
                if not self.t.has_key(keyl):
                    self.ted3(F,G.del_left(),master)
                self.t[(str(F),str(G))] = min(self.t[keyr], self.t[keyl]) + 1
                
            elif master==1 and G.heavy==0 and len(G.roots)>1:
                key = (str(F),str(G.del_right()))
                if not self.t.has_key(key):
                    self.ted3(F,G.del_right(),master)
                self.t[(str(F),str(G))] = self.t[key] + 1
            else:
                key = (str(F),str(G.del_left()))
                if not self.t.has_key(key):
                    self.ted3(F,G.del_left(),master)
                self.t[(str(F),str(G))] = self.t[key] + 1

        elif G.empty():
            if master==1:
                keyr = (str(F.del_right()),str(G))
                if not self.t.has_key(keyr):
                    self.ted3(F.del_right(),G,master)
                keyl = (str(F.del_left()),str(G))
                if not self.t.has_key(keyl):
                    self.ted3(F.del_left(),G,master)
                self.t[(str(F),str(G))] = min(self.t[keyr], self.t[keyl]) + 1

            elif master==0 and F.heavy==0 and len(F.roots)>1:
                key = (str(F.del_right()),str(G))
                if not self.t.has_key(key):
                    self.ted3(F.del_right(),G,master)
                self.t[(str(F),str(G))] = self.t[key] + 1
            else:
                key = (str(F.del_left()),str(G))
                if not self.t.has_key(key):
                    self.ted3(F.del_left(),G,master)                        
                self.t[(str(F),str(G))] = self.t[key] + 1

        else:

            if (master==0 and F.heavy==0 and len(F.roots)>1) or (master==1 and G.heavy==0 and len(G.roots)>1):
                if master==0:                    
                    key0 = (str(F),str(G.del_left()))
                    if not self.t.has_key(key0):
                        self.ted3(F,G.del_left(),master)
                else:
                    key0 = (str(F.del_left()),str(G))
                    if not self.t.has_key(key0):
                        self.ted3(F.del_left(),G,master)

                key1 = (str(F),str(G.del_right()))
                if not self.t.has_key(key1):
                    self.ted3(F,G.del_right(),master)
                key2 = (str(F.del_right()),str(G))
                if not self.t.has_key(key2):
                    self.ted3(F.del_right(),G,master)
                (F3,F4) = F.match_right()
                (G3,G4) = G.match_right()
                key3 = (str(F3),str(G3))
                key4 = (str(F4),str(G4))
                if not self.t.has_key(key3):
                    self.ted3(F3,G3,master)
                if not self.t.has_key(key4):
                    print "problems!!!",key4,key3,key2,key1
                    self.ted3(F4,G4,master)

                self.t[(str(F),str(G))] = min(self.t[key0] + 1 ,
                                              self.t[key1] + 1 ,
                                              self.t[key2] + 1 ,
                                              self.t[key3] + self.t[key4] + (1-(F.roots[-1].label==G.roots[-1].label)))

            else:
                
                if master==0:                    
                    key0 = (str(F),str(G.del_right()))
                    if not self.t.has_key(key0):
                        self.ted3(F,G.del_right(),master)
                else:
                    key0 = (str(F.del_right()),str(G))
                    if not self.t.has_key(key0):
                        self.ted3(F.del_right(),G,master)

                key1 = (str(F),str(G.del_left()))
                if not self.t.has_key(key1):
                    self.ted3(F,G.del_left(),master)
                key2 = (str(F.del_left()),str(G))
                if not self.t.has_key(key2):
                    self.ted3(F.del_left(),G,master)
                (F3,F4) = F.match_left()
                (G3,G4) = G.match_left()
                key3 = (str(F3),str(G3))
                key4 = (str(F4),str(G4))
                if not self.t.has_key(key3):
                    self.ted3(F3,G3,master)
                if not self.t.has_key(key4):
                    print "problems!!!!"
                    self.ted3(F4,G4,master)

                self.t[(str(F),str(G))] = min(self.t[key0] + 1 ,
                                              self.t[key1] + 1 ,
                                              self.t[key2] + 1 ,
                                              self.t[key3] + self.t[key4] + (1-(F.roots[0].label==G.roots[0].label)))

            return self.t[str(F),str(G)]


def post(f):
    '''print left-to-right postorder traversal of forest.'''
    for l in f.depth_first():
        print l


def unrooted(F,G,algo):
    a = 0
    if algo == 'Shasha':
        a = Shasha()
    elif algo == 'Klein':
        a = Klein2()
    else:
        a = DMRW()

    #assume F>G
    nprobs = 0
    mini = 10*F.n
    best=None
    for node in G.depth_first():
        if node!=G.roots[0]:
            g1 = G.change_root(node, node.parent)
            print "now working on ", g1,
            tmp = a.ted(F,g1)
            if (tmp < mini):
                mini = tmp
                best = str(g1)
            print len(a.t) - nprobs
            nprobs = len(a.t)

        for left in node.nbrs:
            g1 = G.change_root(node, left)
            print "now working on ", g1,
            tmp = a.ted(F,g1)
            if (tmp < mini):
                mini = tmp
                best = str(g1)
            print len(a.t) - nprobs
            nprobs = len(a.t)

    return [mini,best,a]
            

def test(n,m):
    f = Forest(sigma=range(5))
    g = Forest(sigma=range(5))
    f.zigzag(n)
    g.zigzag(m)
    s = Shasha()
    k = Klein2()
    d = DMRW()
    print "Shasha: ", s.ted(f,g), "nprobs:",len(s.t)
    k.ted(f,g)
    print "Klein: ", k.ted(f,g), "nprobs:",len(k.t)
    d.ted(f,g)
    print "DMRW: ", d.ted(f,g), "nprobs:",len(d.t)
    return (f,g,s,k,d)


def test2(n,m):
    f = Forest(sigma=range(5))
    g = Forest(sigma=range(5))
#    f.zigzag(n)
#    g.zigzag(m)
    f.balanced(n)
    g.balanced(n)
    s = unrooted(f,g,'Shasha')
    res=[f.n,g.n,len(s[2].t)]
    s = unrooted(f,g,'Klein')
    res.append(len(s[2].t))
    s = unrooted(f,g,'DMRW')
    res.append(len(s[2].t))
    return res

if __name__ == "__main__":
    #a5=Node("e",[],4);
    #a2=Node("b",[a5],1)
    #a3=Node("c",[],2);
    #a4=Node("d",[],3);
    #a1=Node("a",[a2,a3,a4],0)
    #e=Node("e",[],"e");
    d=Node("d",[],"d");
    c=Node("c",[d],"c");
    b=Node("b",[],"b");
    a=Node("a",[b,c],"a");

    a1=Node("a",[],"A");

    f1=Forest([a])
    f2=Forest([a1])

    s = Shasha()
    s.ted(f1,f2)
    for x,y in s.t.items():
        print x,y
    #print str(a1)
    #for tree in a1.depth_first():
    #   print tree


    #foret=Forest([a1])
    #print foret
    #res=foret.match_left()
    #print res[0]
    #print res[1]
    #for tree in foret.depth_first():
    #    print tree



